package cas.ihep.spark.io;

import java.io.*;
import java.nio.ByteBuffer;
import java.util.Formatter;

@SuppressWarnings("WeakerAccess")
public abstract class Packet implements Closeable{
    public abstract int type();
    public abstract int length();
    public abstract int echoSize();
    public abstract int echoReply();
    public abstract int taskCount();
    public abstract int numOfIterations();
    public abstract ByteBuffer getBuffer();
    public abstract <T> T unwrapPacket(Class<T> klass);

    public static final int EchoRequest=0;
    public static final int EchoReply=1;
    public static final int DataRequest=2;
    public static final int DataReply=3;
    public static final int QuitRequest=4;

    public static Packet wrapEchoRequestPacket(int num){
        return PacketImpl.wrapEchoRequestPacket(num);
    }

    public static Packet wrapDataRequestPacket(Object data, int numOfIters){
        return PacketImpl.wrapPacket(DataRequest,data,numOfIters,-1);
    }

    public static Packet wrapDataReplyPacket(Object data, int num, int numOfIters){
        return PacketImpl.wrapPacket(DataReply,data,numOfIters,num);
    }

    public static Packet wrap(ByteBuffer buf){
        return new PacketImpl(buf);
    }

    public static Packet wrap(byte[] buf){
        return new PacketImpl(ByteBuffer.wrap(buf));
    }

    public static Packet QuitPacket(){
        return PacketImpl.PacketHolder.QuitPacket;
    }

    public static Packet EchoReplyOkPacket(){
        return PacketImpl.PacketHolder.EchoReplyOkPacket;
    }

    public static Packet EchoReplyWaitingPacket(){
        return PacketImpl.PacketHolder.EchoReplyWaitingPacket;
    }

    public interface Serializer extends cas.ihep.spark.io.Serializer{}
}

@SuppressWarnings("WeakerAccess")
class PacketByteArrayImpl extends Packet{

    public void close(){}
    static void int2ByteArray(int val,byte[] arr,int off){
        arr[off]= (byte) ((val>>>24)&0xff);
        arr[off+1]= (byte) ((val>>>16)&0xff);
        arr[off+2]= (byte) ((val>>>8)&0xff);
        arr[off+3]= (byte) (val&0xff);
    }
    static int byteArray2Int(byte[] arr, int off){
        return Byte.toUnsignedInt(arr[off])<<24 | Byte.toUnsignedInt(arr[off+1])<<16 | Byte.toUnsignedInt(arr[off+2])<<8 | Byte.toUnsignedInt(arr[off+3]);
    }
    public String toString(){
        StringBuilder builder=new StringBuilder();
        int tp=type();
        new Formatter(builder).format("Packet<offset: %d, size: %d>(type: %s, length: %d",offset,length,type2Str[tp],length());
        if(tp==DataRequest || tp==DataReply){
            builder.append(", iteration: ").append(numOfIterations());
        }
        builder.append(')');
        return builder.toString();
    }

    private PacketByteArrayImpl(){}

    public PacketByteArrayImpl(byte[] arr){
        byteArray=arr;
        offset=0;
        length=arr.length;
    }

    public PacketByteArrayImpl(byte[] arr,int off,int len){
        byteArray=arr;
        reset(off,len);
    }

    private void reset(int off, int len){
        if(len<0){
            throw new IllegalArgumentException("len<0");
        }
        if(off<0 || off+len>byteArray.length){
            throw new IndexOutOfBoundsException();
        }
        offset=off;
        length=len;
    }

    public int type(){
        return byteArray2Int(byteArray,offset+4);
    }

    public int length(){
        return byteArray2Int(byteArray,offset);
    }

    public int echoSize(){
        return byteArray2Int(byteArray,offset+4+4);
    }

    public int echoReply(){return byteArray2Int(byteArray,offset+4+4);}

    public int taskCount(){return byteArray2Int(byteArray,offset+4+4+4);}

    public int numOfIterations(){return byteArray2Int(byteArray,offset+4+4);}

    static Packet wrapInternal(ByteBuffer b){
        if(b.hasArray()) {
            PacketByteArrayImpl p = new PacketByteArrayImpl();
            p.byteArray = b.array();
            p.offset = b.position();
            p.length = b.limit();
            return p;
        }else{
            byte[] arr=new byte[b.remaining()];
            int pos=b.position();
            b.get(arr);
            b.position(pos);
            return new PacketByteArrayImpl(arr);
        }
    }

    public ByteBuffer getBuffer() {
        return ByteBuffer.wrap(byteArray,offset,length);
    }

    private byte[] byteArray;
    private int offset,length;


    static final String[] type2Str={"EchoRequest","EchoReply","DataRequest","DataReply","QuitRequest"};

    static Packet wrapPacket(int type,Object data,int numOfIters,int numOfTasks){
        cas.ihep.spark.io.Serializer s=cas.ihep.spark.io.Serializer.findHandler(data.getClass());
        int totalLength;
        ByteArrayOutputStreamExt bos;
        if(s!=null){
            int len=s.length(data);
            totalLength=len+4+4+4;
            if(type==DataReply){
                totalLength+=4;
            }else{
                numOfTasks=-1;
            }
            bos=new ByteArrayOutputStreamExt(totalLength);
            try(DataOutputStream dos=new DataOutputStream(bos)){
                dos.writeInt(totalLength-4);
                dos.writeInt(type);
                dos.writeInt(numOfIters);
                if(numOfTasks>=0){
                    dos.writeInt(numOfTasks);
                }
                s.serialize(data,dos);
            }catch (IOException ignored){
                throw new RuntimeException("Should never happens");
            }
        }else{
            bos=new ByteArrayOutputStreamExt(32);
            int before;
            DataOutputStream dos=new DataOutputStream(bos);
            try {
                dos.writeInt(0);
                dos.writeInt(type);
                dos.writeInt(numOfIters);
                if(type==DataReply){
                    dos.writeInt(numOfTasks);
                }
            }catch (IOException ignored){}
            before=bos.size();
            try(ObjectOutputStream oos=new ObjectOutputStream(bos)){
                oos.writeObject(data);
            } catch (IOException ignored) {
                throw new RuntimeException("Should never happens");
            }
            int2ByteArray(bos.size()-before,bos.toByteArray(),0);
        }
        return new PacketByteArrayImpl(bos.toByteArray(),0, bos.size());
    }

    private <T> T unwrapPacket1(Class<T> tClass){
        cas.ihep.spark.io.Serializer s=cas.ihep.spark.io.Serializer.findHandler(tClass);
        int dataOffset,dataLength;
        switch (type()){
            case DataRequest:
                dataOffset=offset+12;
                dataLength=length-12;
                break;
            case DataReply:
                dataOffset=offset+16;
                dataLength=length-16;
                break;
            default:
                throw new RuntimeException("Unwrap a packet without data");
        }
        ByteArrayInputStream bis=new ByteArrayInputStream(byteArray,dataOffset,dataLength);
        Object data;
        if(s!=null){
            DataInputStream dis=new DataInputStream(bis);
            data=s.deserialize(dis);
        }else{
            try {
                ObjectInputStream ois=new ObjectInputStream(bis);
                data=ois.readObject();
            } catch (IOException | ClassNotFoundException e) {
                throw new RuntimeException("Should never happens",e);
            }
        }
        return tClass.cast(data);
    }

    public <T> T unwrapPacket(Class<T> t){
        return unwrapPacket1(t);
    }

}
